/*
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package net.javlov;

import net.javlov.vf.OptionValues;
import net.javlov.vf.QFunction;
import net.javlov.vf.Value;

/**
 * Agent implementing the Sarsa algorithm.
 * 
 * @author Matthijs Snel
 *
 */
//TODO! Now this implementation doesn't work with q-functions that are not tabular.
public class SarsaAgent<T> extends TDAgent<T> {

	/**
	 * The Q-value function.
	 */
	protected QFunction<T> q;
	
	/**
	 * The value of the last state-action pair: Q(s_t, a_t).
	 */
	protected double lastQValue;
	
	/**
	 * Last option o_t
	 */
	protected Option<T> lastOption;
	
	/**
	 * Q-values of the current state.
	 */
	protected OptionValues currentQValues;
	
	/**
	 * Default constructor.
	 */
	public SarsaAgent() {}
	
	/**
	 * Constructs agent with given Q-value function and gamma = 0.9.
	 * @param q the value function to use.
	 */
	public SarsaAgent(QFunction<T> q) {
		this(q, 0.9);
	}
	
	/**
	 * Constructs agent with given Q-value function and gamma.
	 * @param q the value function to use.
	 * @param gamma the discountfactor gamma e [0, 1].
	 */
	public SarsaAgent(QFunction<T> q, double gamma) {
		setQFunction(q);
		setGamma(gamma);
	}
	
	public QFunction<T> getQFunction() {
		return q;
	}

	/**
	 * Sets the Q-Function to be the Q-function passed in as argument.
 	 *
	 * @param q the q-function
	 */
	public void setQFunction(QFunction<T> q) {
		this.q = q;
	}
	
	@Override
	public Action doStep( State<? extends T> s, double reward ) {
		currentQValues = q.getValues(s);
		return super.doStep(s, reward);
	}
	
	@Override
	protected Option<T> pickNewOption( State<? extends T> s ) {
		return policy.getOption(s, currentQValues);
	}
	
	@Override
	protected void updateValueFunction(State<? extends T> s, double TDError) {
		learnRate.visit(lastState, lastOption);
		q.update(lastState, lastOption, learnRate.get(lastState,lastOption)*TDError);
		lastQValue = currentQValues.get(currentOption);
		lastState = s;
	}
	
	@Override
	protected double getTDError(State<? extends T> s, double reward) {
		return reward + gamma*currentQValues.get(currentOption) - lastQValue;
	}
	
	@Override
	protected double getSMDPTDError(State<? extends T> s, double reward) {
		return optionAccumulatedReward + optionDiscount*currentQValues.get(currentOption) - lastQValue;
	}
	
	@Override
	public Action<T> firstStep( State<? extends T> s ) {
		currentQValues = q.getValues(s);
		currentOption = pickNewOption(s);
		lastQValue = currentQValues.get(currentOption);
		lastOption = currentOption;
		lastState = s;
		return currentOption.firstStep(s);
	}
		
	@Override
	public void reset() {
		currentOption = null;
		policy.reset();
		lastValue = 0;
		lastState = null;
		lastTDerr = 0;
		lastQValue = 0;
		lastState = null;
		lastOption = null;
		q.reset();
	}
	
	@Override
	public void init() {
		currentOption = null;
		policy.init();
		lastValue = 0;
		lastState = null;
		lastTDerr = 0;
		learnRate.init();
		lastQValue = 0;
		lastState = null;
		lastOption = null;
		q.init();
	}
		
	@Override
	protected void currentOptionHasFinished() {
		currentOptionFinished = true;
		lastOption = currentOption;
	}
	
	/**
	 * Decides whether the option the agent is currently executing should be interrupted. If
	 * {@link #interruptsOptions()} returns false, this method will always return false.
	 * 
	 * Otherwise, an option is interrupted if the value of the state under the current policy
	 * is higher than the Q-value of the current option.
	 * 
	 * @param s the state the agent is in
	 * @return {@code true} if the current option should be interrupted. {@code false} otherwise.
	 */
	@Override
	protected boolean interruptCurrentOption(State<? extends T> s) {
		if ( !interruptOptions )
			return false;
		double 	val = 0,
				probs[] = policy.getOptionProbabilities(s, currentQValues);
		int i = 0;
		for ( Value v : currentQValues.getValues() )
			val += probs[i++]*v.get();
		
		//System.out.println(" V: " + val + ", currQ: " + currentQValues.get(currentOption) + ", allQ: " + currentQValues);
		if ( val > currentQValues.get(currentOption) ) {
			//System.out.println("Interrupting option for " + s);
			//System.out.println(" V: " + val + ", currQ: " + currentQValues.get(currentOption) + ", allQ: " + currentQValues);
			currentOption.setFinished();
			return true;
		}
		return false;
	}
}
